from datetime import datetime
from basic import MongoDBCollectionHandler
from pymongo import ASCENDING


class FollowerCollection(MongoDBCollectionHandler):
    '''
    Class that handles tw_follower collection
    It stores who are following who
    '''
    collection_name = "tw_follower"

    def __init__(self):
        super(FollowerCollection, self).\
                __init__(FollowerCollection.collection_name)
        if(self.col_handler.count() == 0):
            print "Row count is 0. It might be the first time to use " + \
                    FollowerCollection.collection_name + " collection."
            self.col_handler.ensure_index([("screen_name", 1)], unique=True)

    def insert(self, screen_name, follow_list, date=None):
        '''
        Insert new follow_list "ids" of the screen_name for the date.
        The follow_list should be at least empty list not None.
        If the date is not specified, it will save as current time.

        @param screen_name : the screen name to be saved as "leader"
        @param follow_list : the list of the follows of the screen_name.
               It should be at least [], not None
        @param date : save date of the snapshot as the date
        @rtype boolean
        '''

        if (screen_name == None or follow_list == None):
            return False

        if (date == None):
            date = datetime.now()

        #TODO check if the screen_name is in the tw_account collection
        document = {"screen_name": screen_name,
                    "follower": follow_list,
                    "date": date}

        self.col_handler.insert(document)
        return True

    def get_follower(self, screen_name, date=None):
        '''
        Return list of tuple(date, list of follower ids) of the screen name.
        If the date is not None, it will get date until that date.
        If the dats is None, it will return everything!

        @param screen_name : the screen of the account the get follower list
        @param date : get follower list until this date
        '''

        if (screen_name == None):
            return None

        if (date != None):
            follow_cursor = self.col_handler.\
                    find({"screen_name": screen_name,\
                    "date": {"$gte": date}}).sort("date", ASCENDING)
        else:
            follow_cursor = self.col_handler.\
                    find({"screen_name": screen_name}).\
                    sort("date", ASCENDING)

        if (follow_cursor.count() <= 0):
            return None

        follow_list = []
        for index in range(0, follow_cursor.count()):
            info = follow_cursor[index]
            follow_list.append((info["date"], info["follower"]))

        return follow_list


if __name__ == "__main__":
        follower_coll = FollowerCollection()
        #follower_coll.insert("elaborate", [1L, 2L])
        print follower_coll.get_follower("elaborate")
